#include <jni.h>
#include <string>
#include "libyuv.h"

void YV12ToI420(jbyte *src_yv12_data, jint width, jint height, jbyte *dst_i420_data) {
    jbyte *src_yplane = src_yv12_data;
    jbyte *src_uplane = src_yv12_data + width * height;
    jbyte *src_vplane = src_uplane + (width * height / 4);
    jbyte *dst_yplane = dst_i420_data;
    jbyte *dst_uplane = dst_i420_data + width * height;
    jbyte *dst_vplane = dst_uplane + (width * height / 4);
    // YV12 is YVU => Use I420(YUV) copy and flip U and V.
    libyuv::I420Copy(
            (const uint8_t *) src_yplane, width,
            (const uint8_t *) src_vplane, width / 2,
            (const uint8_t *) src_uplane, width / 2,
            (uint8_t *) dst_yplane, width,
            (uint8_t *) dst_uplane, width / 2,
            (uint8_t *) dst_vplane, width / 2,
            width, height);
}

void NV21ToI420(jbyte *src_nv21_data, jint width, jint height, jbyte *dst_i420_data) {
    // NV21 = y plane followed by an interleaved V/U plane, i.e. same as NV12
    // but the U and the V are switched. Use the NV12 function and switch the U
    // and V planes.
    jbyte *src_yplane = src_nv21_data;
    jbyte *src_uvplane = src_nv21_data + width * height;
    jbyte *dst_yplane = dst_i420_data;
    jbyte *dst_uplane = dst_i420_data + width * height;
    jbyte *dst_vplane = dst_uplane + (width * height / 4);
    libyuv::NV21ToI420(
            (const uint8_t *) src_yplane, width,
            (const uint8_t *) src_uvplane, width,
            (uint8_t *) dst_yplane, width,
            (uint8_t *) dst_vplane, width / 2,
            (uint8_t *) dst_uplane, width / 2,
            width, height);
}

void I420ToNv12(jbyte *src_i420_data, jint width, jint height, jbyte *src_nv12_data) {
    jint src_y_size = width * height;
    jint src_u_size = (width >> 1) * (height >> 1);

    jbyte *src_nv12_y_data = src_nv12_data;
    jbyte *src_nv12_uv_data = src_nv12_data + src_y_size;

    jbyte *src_i420_y_data = src_i420_data;
    jbyte *src_i420_u_data = src_i420_data + src_y_size;
    jbyte *src_i420_v_data = src_i420_data + src_y_size + src_u_size;


    libyuv::I420ToNV12(
            (const uint8_t *) src_i420_y_data, width,
            (const uint8_t *) src_i420_u_data, width >> 1,
            (const uint8_t *) src_i420_v_data, width >> 1,
            (uint8_t *) src_nv12_y_data, width,
            (uint8_t *) src_nv12_uv_data, width,
            width, height);
}

void I420ToNv21(jbyte *src_i420_data, jint width, jint height, jbyte *src_nv21_data) {
    jint src_y_size = width * height;
    jint src_u_size = (width >> 1) * (height >> 1);

    jbyte *src_nv12_y_data = src_nv21_data;
    jbyte *src_nv12_uv_data = src_nv21_data + src_y_size;

    jbyte *src_i420_y_data = src_i420_data;
    jbyte *src_i420_u_data = src_i420_data + src_y_size;
    jbyte *src_i420_v_data = src_i420_data + src_y_size + src_u_size;


    libyuv::I420ToNV21(
            (const uint8_t *) src_i420_y_data, width,
            (const uint8_t *) src_i420_u_data, width >> 1,
            (const uint8_t *) src_i420_v_data, width >> 1,
            (uint8_t *) src_nv12_y_data, width,
            (uint8_t *) src_nv12_uv_data, width,
            width, height);
}

void NV21ToNV12(jbyte *src_nv21_data, jbyte *dst_nv12_data, jint width, jint height) {
    jint src_y_size = width * height;

    jbyte *src_nv21_y_data = src_nv21_data;
    jbyte *src_nv21_vu_data = src_nv21_data + src_y_size;

    jbyte *dst_nv12_y_data = dst_nv12_data;
    jbyte *dst_nv12_uv_data = dst_nv12_data + src_y_size;

    libyuv::NV21ToNV12(
            (const uint8_t *) src_nv21_y_data,
            width,
            (const uint8_t *) src_nv21_vu_data,
            width,
            (uint8_t *) dst_nv12_y_data,
            width,
            (uint8_t *) dst_nv12_uv_data,
            width,
            width,
            height
            );
}

void NV21ToNV12Inplace(jbyte *src_nv21_data, jbyte *dst_nv12_data, jint width, jint height) {
    jint src_y_size = width * height;
    jbyte *src_nv21_vu_data = src_nv21_data + src_y_size;
    jbyte *dst_nv12_uv_data = dst_nv12_data;
//NV21ToNV12Inplace(uint8_t *src_vu, int src_stride_vu, uint8_t *tmp_uv, int tmp_stride_uv, int width, int height);
    libyuv::NV21ToNV12Inplace(
            (uint8_t *) src_nv21_vu_data,
            width,
            (uint8_t *) dst_nv12_uv_data,
            width,
            width,
            height
            );
}

void rotateI420(jbyte *src_i420_data, jint width, jint height, jbyte *dst_i420_data, jint degree) {
    jint src_i420_y_size = width * height;
    jint src_i420_u_size = (width >> 1) * (height >> 1);

    jbyte *src_i420_y_data = src_i420_data;
    jbyte *src_i420_u_data = src_i420_data + src_i420_y_size;
    jbyte *src_i420_v_data = src_i420_data + src_i420_y_size + src_i420_u_size;

    jbyte *dst_i420_y_data = dst_i420_data;
    jbyte *dst_i420_u_data = dst_i420_data + src_i420_y_size;
    jbyte *dst_i420_v_data = dst_i420_data + src_i420_y_size + src_i420_u_size;

    //要注意这里的width和height在旋转之后是相反的
    if (degree == libyuv::kRotate90 || degree == libyuv::kRotate270) {
        libyuv::I420Rotate((const uint8_t *) src_i420_y_data, width,
                           (const uint8_t *) src_i420_u_data, width >> 1,
                           (const uint8_t *) src_i420_v_data, width >> 1,
                           (uint8_t *) dst_i420_y_data, height,
                           (uint8_t *) dst_i420_u_data, height >> 1,
                           (uint8_t *) dst_i420_v_data, height >> 1,
                           width, height,
                           (libyuv::RotationMode) degree);
    } else {
        libyuv::I420Rotate((const uint8_t *) src_i420_y_data, width,
                           (const uint8_t *) src_i420_u_data, width >> 1,
                           (const uint8_t *) src_i420_v_data, width >> 1,
                           (uint8_t *) dst_i420_y_data, width,
                           (uint8_t *) dst_i420_u_data, width >> 1,
                           (uint8_t *) dst_i420_v_data, width >> 1,
                           width, height,
                           (libyuv::RotationMode) degree);
    }
}

extern "C"
JNIEXPORT void JNICALL
Java_org_doubango_jni_Libyuv_YV12ToI420(JNIEnv *env, jclass clazz, jbyteArray src, jbyteArray dest,
                                        jint width, jint height) {
    jbyte *srcP = env->GetByteArrayElements(src, NULL);
    jbyte *destP = env->GetByteArrayElements(dest, NULL);
    YV12ToI420(srcP, width, height, destP);
    env->ReleaseByteArrayElements(src, srcP, 0);
    env->ReleaseByteArrayElements(dest, destP, 0);
}

extern "C"
JNIEXPORT void JNICALL
Java_org_doubango_jni_Libyuv_NV21ToI420(JNIEnv *env, jclass clazz, jbyteArray nv21_bytes,
                                        jbyteArray i420_bytes, jint width, jint height) {
    jbyte *nv21_bytes_p = env->GetByteArrayElements(nv21_bytes, NULL);
    jbyte *i420_bytes_p = env->GetByteArrayElements(i420_bytes, NULL);
    // nv21转化为i420
    NV21ToI420(nv21_bytes_p, width, height, i420_bytes_p);
    env->ReleaseByteArrayElements(nv21_bytes, nv21_bytes_p, 0);
    env->ReleaseByteArrayElements(i420_bytes, i420_bytes_p, 0);
}

extern "C"
JNIEXPORT void JNICALL
Java_org_doubango_jni_Libyuv_I420Rotate(JNIEnv *env, jclass clazz, jbyteArray src_i420_bytes,
                                        jbyteArray dest_i420_bytes, jint width, jint height,
                                        jint degree) {
    jbyte *src_i420_bytes_p = env->GetByteArrayElements(src_i420_bytes, NULL);
    jbyte *dest_i420_bytes_p = env->GetByteArrayElements(dest_i420_bytes, NULL);
    rotateI420(src_i420_bytes_p, width, height, dest_i420_bytes_p, degree);
    env->ReleaseByteArrayElements(src_i420_bytes, src_i420_bytes_p, 0);
    env->ReleaseByteArrayElements(dest_i420_bytes, dest_i420_bytes_p, 0);
}
extern "C"
JNIEXPORT void JNICALL
Java_org_doubango_jni_Libyuv_I420ToNV12(JNIEnv *env, jclass clazz, jbyteArray i420_bytes,
                                        jbyteArray nv12_bytes, jint width, jint height) {
    jbyte *i420_bytes_p = env->GetByteArrayElements(i420_bytes, NULL);
    jbyte *nv12_bytes_p = env->GetByteArrayElements(nv12_bytes, NULL);
    I420ToNv12(i420_bytes_p, width, height, nv12_bytes_p);
    env->ReleaseByteArrayElements(i420_bytes, i420_bytes_p, 0);
    env->ReleaseByteArrayElements(nv12_bytes, nv12_bytes_p, 0);
}

extern "C"
JNIEXPORT void JNICALL
Java_org_doubango_jni_Libyuv_YV12ToNV12Rotate180(JNIEnv *env, jclass clazz, jbyteArray yv12_bytes,
                                                 jbyteArray nv12_bytes, jint width, jint height) {
    jbyte *yv12_bytes_p = env->GetByteArrayElements(yv12_bytes, NULL);
    jbyte *nv12_bytes_p = env->GetByteArrayElements(nv12_bytes, NULL);
//    const uint8_t *src_y, int src_stride_y, const uint8_t *src_v,
//    int src_stride_v, const uint8_t *src_u, int src_stride_u,
//    uint8_t *dst_y, int dst_stride_y, uint8_t *dst_uv,
//            int dst_stride_uv, int width, int height

    jint src_yv12_y_size = width * height;
    jint src_yv12_v_size = (width >> 1) * (height >> 1);

    jbyte *src_yv12_y_data = yv12_bytes_p;
    jbyte *src_yv12_v_data = yv12_bytes_p + src_yv12_y_size;
    jbyte *src_yv12_u_data = yv12_bytes_p + src_yv12_y_size + src_yv12_v_size;

    jbyte *dst_nv12_y_data = nv12_bytes_p;
    jbyte *dst_nv12_uv_data = nv12_bytes_p + src_yv12_y_size;


    libyuv::YV12ToNV12Rotate180(reinterpret_cast<const uint8_t *>(src_yv12_y_data), width,
                                reinterpret_cast<const uint8_t *>(src_yv12_v_data), width >> 1,
                                reinterpret_cast<const uint8_t *>(src_yv12_u_data), width >> 1,
                                reinterpret_cast<uint8_t *>(dst_nv12_y_data), width,
                                reinterpret_cast<uint8_t *>(dst_nv12_uv_data), width, width, height);
    env->ReleaseByteArrayElements(yv12_bytes, yv12_bytes_p, 0);
    env->ReleaseByteArrayElements(nv12_bytes, nv12_bytes_p, 0);
}
extern "C"
JNIEXPORT void JNICALL
Java_org_doubango_jni_Libyuv_YV12ToNV12Rotate270Or90(JNIEnv *env, jclass clazz, jbyteArray yv12_bytes,
                                                 jbyteArray nv12_bytes, jint width, jint height, jboolean is270) {
    jbyte *yv12_bytes_p = env->GetByteArrayElements(yv12_bytes, NULL);
    jbyte *nv12_bytes_p = env->GetByteArrayElements(nv12_bytes, NULL);
//    const uint8_t *src_y, int src_stride_y, const uint8_t *src_v,
//    int src_stride_v, const uint8_t *src_u, int src_stride_u,
//    uint8_t *dst_y, int dst_stride_y, uint8_t *dst_uv,
//            int dst_stride_uv, int width, int height

    jint src_yv12_y_size = width * height;
    jint src_yv12_v_size = (width >> 1) * (height >> 1);

    jbyte *src_yv12_y_data = yv12_bytes_p;
    jbyte *src_yv12_v_data = yv12_bytes_p + src_yv12_y_size;
    jbyte *src_yv12_u_data = yv12_bytes_p + src_yv12_y_size + src_yv12_v_size;

    jbyte *dst_nv12_y_data = nv12_bytes_p;
    jbyte *dst_nv12_uv_data = nv12_bytes_p + src_yv12_y_size;


    libyuv::YV12ToNV12Rotate270Or90(reinterpret_cast<const uint8_t *>(src_yv12_y_data), width,
                                reinterpret_cast<const uint8_t *>(src_yv12_v_data), width >> 1,
                                reinterpret_cast<const uint8_t *>(src_yv12_u_data), width >> 1,
                                reinterpret_cast<uint8_t *>(dst_nv12_y_data), height,
                                reinterpret_cast<uint8_t *>(dst_nv12_uv_data), height, width, height, is270);
    env->ReleaseByteArrayElements(yv12_bytes, yv12_bytes_p, 0);
    env->ReleaseByteArrayElements(nv12_bytes, nv12_bytes_p, 0);
}

extern "C"
JNIEXPORT void JNICALL
Java_org_doubango_jni_Libyuv_NV21ToNV12(JNIEnv *env, jclass clazz, jbyteArray nv21_bytes,
                                        jbyteArray nv12_bytes, jint width, jint height) {
    jbyte *nv21_bytes_p = env->GetByteArrayElements(nv21_bytes, NULL);
    jbyte *nv12_bytes_p = env->GetByteArrayElements(nv12_bytes, NULL);
    NV21ToNV12(nv21_bytes_p, nv12_bytes_p, width, height);
    env->ReleaseByteArrayElements(nv21_bytes, nv21_bytes_p, 0);
    env->ReleaseByteArrayElements(nv12_bytes, nv12_bytes_p, 0);
}
extern "C"
JNIEXPORT void JNICALL
Java_org_doubango_jni_Libyuv_NV21ToNV12Inplace(JNIEnv *env, jclass clazz, jbyteArray nv21_bytes,
                                               jbyteArray nv12_bytes, jint width, jint height) {
    jbyte *nv21_bytes_p = env->GetByteArrayElements(nv21_bytes, NULL);
    jbyte *nv12_bytes_p = env->GetByteArrayElements(nv12_bytes, NULL);
    NV21ToNV12Inplace(nv21_bytes_p, nv12_bytes_p, width, height);
    env->ReleaseByteArrayElements(nv21_bytes, nv21_bytes_p, 0);
    env->ReleaseByteArrayElements(nv12_bytes, nv12_bytes_p, 0);
}
extern "C"
JNIEXPORT void JNICALL
Java_org_doubango_jni_Libyuv_I420ToNV21(JNIEnv *env, jclass clazz, jbyteArray i420_bytes,
                                        jbyteArray nv21_bytes, jint width, jint height) {
    jbyte *i420_bytes_p = env->GetByteArrayElements(i420_bytes, NULL);
    jbyte *nv21_bytes_p = env->GetByteArrayElements(nv21_bytes, NULL);
    I420ToNv21(i420_bytes_p, width, height, nv21_bytes_p);
    env->ReleaseByteArrayElements(i420_bytes, i420_bytes_p, 0);
    env->ReleaseByteArrayElements(nv21_bytes, nv21_bytes_p, 0);
}